# 并行链实现产品评论
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda, RunnableParallel
from model_utils import getLLM

model = getLLM()

prompt_template = ChatPromptTemplate.from_messages([
    ("system", "你是一个专业的产品评论员"),
    ("human", "列出产品的主要特点{product_name}"),
])

advantage_template = ChatPromptTemplate.from_messages([
    ("system", "你是专业的产品评论员"),
    ("human", "鉴于这些功能：{features},列出这些功能的优点"),
])

disadvantage_template = ChatPromptTemplate.from_messages([
    ("system", "你是专业的产品评论员"),
    ("human", "鉴于这些功能：{features},列出这些功能的缺点"),
])

# 定义并行执行链
def make_chain():
    # 先获取features
    features_chain = prompt_template | model | StrOutputParser()

    # 根据features获取优点和缺点
    advantage_disadvantage_chain = RunnableParallel(
        pros = advantage_template | model | StrOutputParser(),
        cons = disadvantage_template | model | StrOutputParser(),
    )

    # 将features的输出传给advantage_disadvantage_chain
    full_chain = features_chain | RunnableLambda(
        lambda features:advantage_disadvantage_chain.invoke({"features": features})
    )

    return full_chain

# 创建chain
chain = make_chain()

# 执行
result = chain.invoke({"product_name": "奥迪汽车"})

# 输出结果
print(result)
